import mcp_client
import json
from datetime import datetime

def find_available_tickets_for_train(train_code, departure_cities, arrival_cities, query_date=None):
    """
    使用 mcp-server-chinarailway 服务循环遍历站点组合，查询指定车次在特定日期的余票情况。

    参数:
    - train_code (str): 列车车次，例如 "G28"
    - departure_cities (list): 出发城市名称列表，例如 ["北京南", "济南西"]
    - arrival_cities (list): 到达城市名称列表，例如 ["南京南", "上海虹桥"]
    - query_date (str, optional): 查询日期，格式 "YYYY-MM-DD"。如果为 None，则默认为今天。
    
    返回:
    - list: 一个包含有票信息的字典列表
    """
    
    # 1. 初始化 MCP 客户端
    try:
        client = mcp_client.Client()
    except Exception as e:
        print(f"初始化 MCP 客户端失败，请检查网络或依赖: {e}")
        return []

    # 2. 定义席位代码到中文名称的映射
    seat_map = {
        'swz_num': '商务座',
        'zy_num': '一等座',
        'ze_num': '二等座',
        'gr_num': '高级软卧',
        'rw_num': '软卧',
        'dw_num': '动卧',
        'yw_num': '硬卧',
        'rz_num': '软座',
        'yz_num': '硬座',
        'wz_num': '无座'
    }

    # 3. 如果未提供日期，则默认为今天
    if not query_date:
        query_date = datetime.now().strftime('%Y-%m-%d')
    
    available_tickets_list = []

    print(f"任务开始：查询车次 【{train_code}】 在日期 【{query_date}】 的余票情况。\n")

    # 4. 循环遍历所有出发-到达组合
    for dep_city in departure_cities:
        for arr_city in arrival_cities:
            if dep_city == arr_city:
                continue

            print(f"正在查询: {dep_city} -> {arr_city}...")

            try:
                # 5. 调用 MCP API，传入指定的日期
                api_result = client.call(
                    "@shenpeiheng/mcp-server-chinarailway", 
                    "queryTickets", 
                    {
                        "train_date": query_date, # 使用传入的日期参数
                        "from_station": dep_city,
                        "to_station": arr_city
                    }
                )

                # 6. 解析返回结果，查找指定车次
                if not api_result:
                    print(f"  -> API 未返回有效数据。")
                    continue

                found_train = False
                for train_info in api_result:
                    if train_info.get("station_train_code") == train_code:
                        found_train = True
                        available_seats = {}
                        
                        for seat_code, seat_name in seat_map.items():
                            seat_status = train_info.get(seat_code)
                            if seat_status and seat_status not in ["无", "--", "0", ""]:
                                available_seats[seat_name] = seat_status
                        
                        if available_seats:
                            result = {
                                "日期": query_date,
                                "车次": train_code,
                                "出发地": dep_city,
                                "到达地": arr_city,
                                "余票席位": available_seats
                            }
                            available_tickets_list.append(result)
                            print(f"  --> ✅ 查询成功！发现余票: {available_seats}")
                        else:
                            print(f"  --> 🟡 已找到车次 {train_code}，但当前无余票。")
                        
                        break
                
                if not found_train:
                    print(f"  --> ❌ 在 {dep_city} 到 {arr_city} 的线路中未找到车次 {train_code}。")

            except Exception as e:
                print(f"  --> 🚨 调用 API 查询 {dep_city} 到 {arr_city} 时发生错误: {e}")

    return available_tickets_list

# --- 使用示例 ---
if __name__ == '__main__':
    # ==================== 请在这里修改您要查询的信息 ====================
    # 1. 输入您要查询的日期，格式为 "YYYY-MM-DD"
    #    如果留空字符串 ""，则默认查询今天
    target_date = "2025-09-24" 

    # 2. 输入您要查询的单一车次
    target_train = "G2694"  # 示例车次：G1 (北京南 -> 上海虹桥)

    # 3. 输入所有可能的出发站点
    departure_points = ["深圳北站"]

    # 4. 输入所有可能的到达站点
    arrival_points = ["太湖南站"]
    # =================================================================

    # 执行查询函数
    results = find_available_tickets_for_train(target_train, departure_points, arrival_points, target_date)

    # 打印最终汇总结果
    print("\n---------- ✅ 查询任务完成，结果汇总如下 ----------")
    if results:
        # 使用json.dumps美化输出
        print(json.dumps(results, ensure_ascii=False, indent=4))
    else:
        print("所有查询组合均未发现符合条件的余票。")